import cv2
import numpy as np
import mediapipe as mp
import os
import time
from datetime import datetime

class CleanFaceMeshDetector:
    def __init__(self, output_dir="clean_face_keypoints"):
        """
        初始化纯净显示的人脸关键点检测器
        :param output_dir: 保存关键点数据的目录
        """
        self.output_dir = output_dir
        self.cap = None
        self.collected_count = 0
        
        # 初始化MediaPipe
        self.mp_face_mesh = mp.solutions.face_mesh
        self.mp_drawing = mp.solutions.drawing_utils
        self.mp_drawing_styles = mp.solutions.drawing_styles
        
        # 创建输出目录
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            print(f"创建目录: {output_dir}")
    
    def start_camera(self):
        """启动摄像头"""
        self.cap = cv2.VideoCapture(0)
        if not self.cap.isOpened():
            print("无法打开摄像头")
            return False
        print("摄像头启动成功")
        return True
    
    def detect_face_mesh_clean(self):
        """
        检测人脸关键点并显示纯净的动态效果（无文字信息）
        """
        if self.cap is None:
            print("请先启动摄像头")
            return
        
        print("按 'c' 键保存当前关键点数据")
        print("按 'q' 键退出程序")
        print("按 's' 键开启/关闭自动保存模式")
        
        auto_save = False
        last_save_time = 0
        save_interval = 3  # 自动保存间隔（秒）
        
        # 初始化MediaPipe Face Mesh
        with self.mp_face_mesh.FaceMesh(
            max_num_faces=1,
            refine_landmarks=True,
            min_detection_confidence=0.5,
            min_tracking_confidence=0.5
        ) as face_mesh:
            
            while True:
                ret, frame = self.cap.read()
                if not ret:
                    print("无法读取摄像头画面")
                    break
                
                # 转换为RGB格式（MediaPipe需要RGB）
                rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                
                # 检测人脸关键点
                results = face_mesh.process(rgb_frame)
                
                # 创建两个纯净的显示窗口（无文字信息）
                original_with_mesh = frame.copy()
                mesh_only = np.zeros_like(frame)
                
                if results.multi_face_landmarks:
                    for face_landmarks in results.multi_face_landmarks:
                        # 在原始图像上绘制网格（纯净版本）
                        self.mp_drawing.draw_landmarks(
                            image=original_with_mesh,
                            landmark_list=face_landmarks,
                            connections=self.mp_face_mesh.FACEMESH_TESSELATION,
                            landmark_drawing_spec=None,
                            connection_drawing_spec=self.mp_drawing_styles.get_default_face_mesh_tesselation_style()
                        )
                        
                        # 在黑色背景上绘制网格（纯净版本）
                        self.mp_drawing.draw_landmarks(
                            image=mesh_only,
                            landmark_list=face_landmarks,
                            connections=self.mp_face_mesh.FACEMESH_TESSELATION,
                            landmark_drawing_spec=None,
                            connection_drawing_spec=self.mp_drawing_styles.get_default_face_mesh_tesselation_style()
                        )
                
                # 显示两个纯净窗口（无任何文字信息）
                cv2.imshow('Original with Face Mesh', original_with_mesh)
                cv2.imshow('Face Mesh Only (Black Background)', mesh_only)
                
                # 键盘事件处理
                key = cv2.waitKey(1) & 0xFF
                
                if key == ord('q'):
                    break
                elif key == ord('c'):
                    self._save_clean_keypoints_data(results.multi_face_landmarks, frame)
                elif key == ord('s'):
                    auto_save = not auto_save
                    print(f"自动保存模式: {'开启' if auto_save else '关闭'}")
                
                # 自动保存模式
                if auto_save and results.multi_face_landmarks:
                    current_time = time.time()
                    if current_time - last_save_time >= save_interval:
                        self._save_clean_keypoints_data(results.multi_face_landmarks, frame)
                        last_save_time = current_time
    
    def _save_clean_keypoints_data(self, face_landmarks_list, frame):
        """
        保存纯净的关键点数据（无文字信息）
        :param face_landmarks_list: 人脸关键点列表
        :param frame: 原始图像帧
        """
        if not face_landmarks_list:
            print("未检测到人脸关键点")
            return
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        for i, face_landmarks in enumerate(face_landmarks_list):
            # 提取关键点坐标
            landmarks_data = []
            for landmark in face_landmarks.landmark:
                landmarks_data.append({
                    'x': landmark.x,
                    'y': landmark.y,
                    'z': landmark.z
                })
            
            # 保存关键点数据为JSON格式
            import json
            data_filename = f"clean_keypoints_{timestamp}_{i+1}.json"
            data_filepath = os.path.join(self.output_dir, data_filename)
            
            with open(data_filepath, 'w', encoding='utf-8') as f:
                json.dump({
                    'timestamp': timestamp,
                    'landmarks_count': len(landmarks_data),
                    'landmarks': landmarks_data
                }, f, indent=2, ensure_ascii=False)
            
            # 保存纯净的带网格图像（无文字信息）
            img_filename = f"clean_face_mesh_{timestamp}_{i+1}.jpg"
            img_filepath = os.path.join(self.output_dir, img_filename)
            
            # 创建纯净的带网格图像用于保存
            clean_mesh_image = frame.copy()
            self.mp_drawing.draw_landmarks(
                image=clean_mesh_image,
                landmark_list=face_landmarks,
                connections=self.mp_face_mesh.FACEMESH_TESSELATION,
                landmark_drawing_spec=None,
                connection_drawing_spec=self.mp_drawing_styles.get_default_face_mesh_tesselation_style()
            )
            
            cv2.imwrite(img_filepath, clean_mesh_image)
            self.collected_count += 1
            print(f"已保存纯净关键点数据: {data_filename} 和图像: {img_filename}")
    
    def release(self):
        """释放资源"""
        if self.cap:
            self.cap.release()
        cv2.destroyAllWindows()
        print(f"程序结束，共保存了 {self.collected_count} 组纯净关键点数据")

def main():
    """主函数"""
    print("=== 纯净人脸关键点检测程序 ===")
    print("功能说明:")
    print("1. 程序会实时检测人脸关键点并显示绿色网格")
    print("2. 左侧窗口显示原始图像+网格（无文字）")
    print("3. 右侧窗口显示纯网格效果（黑色背景，无文字）")
    print("4. 按 'c' 键手动保存当前关键点数据")
    print("5. 按 's' 键开启/关闭自动保存模式")
    print("6. 按 'q' 键退出程序")
    print("=" * 40)
    
    # 创建纯净人脸关键点检测器
    detector = CleanFaceMeshDetector()
    
    # 启动摄像头
    if not detector.start_camera():
        return
    
    try:
        # 开始检测和显示
        detector.detect_face_mesh_clean()
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    finally:
        # 释放资源
        detector.release()

if __name__ == "__main__":
    main() 